'''
难度 简单
题目链接 https://leetcode.cn/problems/take-gifts-from-the-richest-pile/description/
'''

from typing import List
import heapq, math

class Solution:
    # 使用内置大顶堆
    def pickGifts1(self, gifts: List[int], k: int) -> int:
        n = len(gifts)
        # 初始化堆大顶堆
        gifts = heapq.nlargest(n, gifts)
        # print(gifts)

        # 取值操作 堆化
        for _ in range(k):
            gifts[0] = math.floor(math.sqrt(gifts[0]))
            gifts = heapq.nlargest(n, gifts)

        return sum(gifts)
    
    # 使用内置大顶堆 优化一下 不需要每次重新建堆
    def pickGifts(self, gifts: List[int], k: int) -> int:
        # 因为默认是小顶堆 所以我初始化取负数 当做大顶堆来用
        heap = [-gift for gift in gifts]
        # 堆化
        heapq.heapify(heap)
        # print(heap)

        # 取值操作 堆化
        for _ in range(k):
            x = heapq.heappop(heap)
            heapq.heappush(heap, -int(math.sqrt(-x)))

        return -sum(heap)

if __name__ == '__main__':
    s = Solution()
    参数 = [
        [[25,64,9,4,100], 4],
        [[1,1,1,1], 4],
    ]
    预期 = [
        29,
        4,
    ]
    for k,v in enumerate(参数):
        r = s.pickGifts(v[0], v[1])
        if r != 预期[k]:
            print("第{}个case不通过 期望={} 结果={}\n".format(
                k+1,
                预期[k],
                r
            ))